package com.stanford.ml.naivebayes;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;

import com.stanford.data.User;

/**
 * Naive Bayes classifier is a simple probabilistic classifier based on 
 * applying Bayes' theorem with strong (naive) independence assumptions.
 * 
 * @author Fatih Sunor
 */
public class NaiveBayes {
	
	private TrainingSet trainingSet;
	private Map<ClassLabel,Double> classPriors;
	private Map<ClassLabel,Map<Attribute, Data>> probabilities;
	private ArrayList<String> attributeList;
	private Map<String, Integer> tagMax;
	
	/**
	 * Constructor that defaults classifier details
	 * @param set is the training set
	 * @param users is the list of users
	 */
	public NaiveBayes(TrainingSet set, List<User> users) {
		trainingSet = set;
		tagMax = populateTagMax(users);
		classPriors = new HashMap<ClassLabel,Double>(trainingSet.getNumberOfClassLabels());
	}

	private Map<String, Integer> populateTagMax(List<User> users) {
		Map<String, Integer> maxTag = new HashMap<String, Integer>();
		for(User user: users){
			Map<String, Integer> userTags = user.getTags();
			for(String tag: userTags.keySet()){
				if(maxTag.containsKey(tag)){
					if(userTags.get(tag)>maxTag.get(tag)){
						maxTag.put(tag, userTags.get(tag));
					}
				}else{
					maxTag.put(tag, userTags.get(tag));
				}
			}
		}
		return maxTag;
	}

	/**
	 * Classifies a given feature set
	 * @param instance is the data with feature set
	 * @return 1 or -1 rating
	 */
	public ClassLabel classify(Instance instance) {
		ClassLabel bestConcept = null;
		double bestP = 0.0;
		for (ClassLabel c : trainingSet.classLabels()) {
			double p = getProbability(c, instance);
			if( p >= bestP ){
				bestConcept = c;
				bestP = p;
			}
		}
		return bestConcept;
	}

	/**
	 * Trains the algorithm
	 */
	public void train() {
		calculateConceptPriors();
		calculateConditionalProbabilities();
	}

	/**
	 * Attaches a particular attribute to the model to be trained
	 * @param name
	 */
	public void trainOnAttribute(String name) {
		if (attributeList ==null) {
			attributeList = new ArrayList<String>();
		}
		attributeList.add(name);
	}
	
	private void calculateConceptPriors() {
		for (ClassLabel c : trainingSet.classLabels()) {
			int totalConceptCount=0;
			for (Instance i : trainingSet.getInstances().values()) {
				if (i.getClassLabel().equals(c)) {
					totalConceptCount++;
				}
			}
			classPriors.put(c, new Double(totalConceptCount));
		}
	}
	
	private void calculateConditionalProbabilities() {
		probabilities = new HashMap<ClassLabel, Map<Attribute, Data>>();
		for (Instance example : trainingSet.getInstances().values()) {
			for (Attribute attribute: example.getAttributes(tagMax)) {
				if (attribute != null && attributeList.contains(attribute.getName())) {
					if (probabilities.get(example.getClassLabel())== null ) {
						probabilities.put(example.getClassLabel(), new HashMap<Attribute, Data>());
					}
					Map<Attribute, Data> probabilityMap = probabilities.get(example.getClassLabel());
					Data data = probabilityMap.get(attribute);
					if (data == null) {
					  data = new Data(attribute.getValue()); 
						probabilityMap.put(attribute, data);
					} else {
						data.increment();
					}
				}
			}
		}		
	}
	
	/**
	 * Calculates the class posterior probabilities
	 * @param classLabel is the outcome rating
	 * @param instance is the data
	 * @return posterior probability
	 */
	public double getProbability(ClassLabel classLabel, Instance instance) {
		return (getProbability(instance,classLabel)*getProbability(classLabel))/getProbability(instance);
	}

	/**
	 * Calculates the probability of an instance
	 * @param instance
	 * @return the probability of an instances
	 */
	public double getProbability(Instance instance) {
		double classPriors=0;
		for (ClassLabel c : getTrainingSet().classLabels()) {
			classPriors += getProbability(instance,c)*getProbability(c);
		}
		return (classPriors == 0) ? (double)1/trainingSet.getSize() : classPriors;
	}
	
	/**
	 * Calculates the probability for a class label
	 * @param classLabel
	 * @return
	 */
	public double getProbability(ClassLabel classLabel) {
		Double instanceCount = classPriors.get(classLabel);
		if(instanceCount == null){
			instanceCount = 0.0;
		}
		return instanceCount/trainingSet.getSize();	
	}
	
	/**
	 * Calculates conditional class probability
	 * @param instance is the data
	 * @param classLabel is the class label
	 * @return a double probability value
	 */
	public double getProbability(Instance instance, ClassLabel classLabel) {
		double classProbability = 1;
		for (Attribute attribute : instance.getAttributes(tagMax)) {
			if ( attribute != null && attributeList.contains(attribute.getName()) ) {
				Map<Attribute, Data> aMap = probabilities.get(classLabel);
				Data data = aMap.get(attribute); 
				if ( data == null) {
					classProbability *= ((double) 1 / (trainingSet.getSize()+1));
				} else {
				  classProbability *= (data.getCount()/classPriors.get(classLabel));
				}
			}
		}	
		return (classProbability == 1) ? (double)1/trainingSet.getNumberOfClassLabels() : classProbability;
	}

	/**
	 * @return the training set
	 */
	public TrainingSet getTrainingSet() {
		return trainingSet;
	}

	/**
	 * Max occurances is necessary to normalize the tag counts
	 * @return the tags with maximum occurances
	 */
	public Map<String, Integer> getTagMaxMap() {
		return tagMax;
	}

	/**
	 * Indicate which attributes are to be trained
	 * @param allTags
	 */
	public void trainOnAttribute(Set<String> allTags) {
		for(String tag: allTags){
			trainOnAttribute(tag);
		}
	}	
}
